package com.whz.gateway.common.util;

import com.google.common.collect.Maps;
import com.whz.generic.Constants;
import com.whz.governance.http.ribbon.context.RibbonFilterContextHolder;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpCookie;
import org.springframework.http.ResponseCookie;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;

import java.util.Map;

/**
 * 灰度测试支持
 *
 * @author whz
 */
@Slf4j
@UtilityClass
public class GreyUtils {
    public static String greyTestRibbon(ServerRequest request) {
        String greyTag = getFlag(request);
        saveGreyFlag(greyTag);
        return greyTag;
    }

    public static String greyTestRibbon(ServerWebExchange exchange) {
        String greyTag = getGreyTag(exchange);
        saveGreyFlag(greyTag);
        return greyTag;
    }


    private static String getFlag(final ServerRequest request) {
        String greyTag = request.headers().firstHeader(Constants.GREY_HEAD_KEY);
        if (StringUtils.hasText(greyTag)) {
            return greyTag;
        }
        final HttpCookie cookie = request.cookies().getFirst(Constants.GREY_HEAD_KEY);
        if (cookie != null) {
            greyTag = cookie.getValue();
        }
        return greyTag;
    }

    private static String getGreyTag(final ServerWebExchange exchange) {
        String greyTag = exchange.getRequest().getHeaders().getFirst(Constants.GREY_HEAD_KEY);
        if (StringUtils.hasText(greyTag)) {
            return greyTag;
        }
        final HttpCookie cookie = exchange.getRequest().getCookies().getFirst(Constants.GREY_HEAD_KEY);
        if (cookie != null) {
            greyTag = cookie.getValue();
        }
        return greyTag;
    }

    private static void saveGreyFlag(final String greyTag) {
        if (!StringUtils.isEmpty(greyTag)) {
            log.info("header`s grey-flag is [{}]", greyTag);
            // 开启ribbon的灰度开关（线程级别）
            RibbonFilterContextHolder.getCurrentContext().add(Constants.GREY_FLAG, greyTag);
        } else {
            RibbonFilterContextHolder.getCurrentContext().remove(Constants.GREY_FLAG);
        }
    }

    public static ServerRequest addGreyFlag(ServerRequest request, String greyFlag) {
        if (StringUtils.hasText(greyFlag) && StringUtils.isEmpty(request.headers().firstHeader(Constants.GREY_HEAD_KEY))) {
            return ServerRequest.from(request)
                    .header(Constants.GREY_HEAD_KEY, greyFlag)
                    .body(request.exchange().getRequest().getBody())
                    .build();
        } else {
            return request;
        }
    }

    public static ServerWebExchange addGreyFlag(ServerWebExchange exchange, String greyFlag) {
        if (StringUtils.hasText(greyFlag) && StringUtils.isEmpty(exchange.getRequest().getHeaders().getFirst(Constants.GREY_HEAD_KEY))) {
            Map<String, String> header = Maps.newHashMap();
            header.put(Constants.GREY_HEAD_KEY, greyFlag);
            return HeaderUtils.addIfAbsent(exchange, header);
        } else {
            return exchange;
        }
    }

    public ResponseCookie buildCookie(final String greyFlag) {
        if (StringUtils.hasText(greyFlag)) {
            return ResponseCookie.from(Constants.GREY_HEAD_KEY, greyFlag)
                    //.secure(true)
                    .httpOnly(true)
                    .maxAge(-1)
                    .path("/")
                    .build();
        }
        return null;
    }
}